import numpy as np
import matplotlib.pyplot as plt
from brian2 import *
import networkx as nx
from matplotlib.animation import FuncAnimation
import matplotlib.animation as animation

# 设置Brian2的单位和后端
prefs.codegen.target = 'numpy'  # 显式设置numpy后端
defaultclock.dt = 0.1*ms

# 创建神经元群
N = 100  # 神经元数量
tau = 10*ms  # 时间常数
eqs = '''
dv/dt = (v0 - v)/tau : volt
v0 : volt
'''

# 创建神经元群组
neurons = NeuronGroup(N, eqs,
                     threshold='v>15*mV',
                     reset='v=0*mV',
                     method='exact')
neurons.v0 = '20*mV * i / (N-1)'

# 创建突触连接
S = Synapses(neurons, neurons, 'w : volt', on_pre='v_post += w')
S.connect(condition='i!=j', p=0.1)  # 10%的连接概率
S.w = '0.5*mV'

# 记录数据
M = StateMonitor(neurons, 'v', record=True)
spikemon = SpikeMonitor(neurons)

# 运行模拟
duration = 500*ms
run(duration)

# 将数据转换为numpy数组以便于处理
t_points = np.array(M.t/ms)
v_points = np.array(M.v/mV)
spike_times = np.array(spikemon.t/ms)
spike_indices = np.array(spikemon.i)

# 创建动画
fig = plt.figure(figsize=(15, 10))

# 准备网络布局
G = nx.DiGraph()
for i in range(N):
    G.add_node(i)
for i, j in zip(S.i, S.j):
    G.add_edge(i, j)
pos = nx.spring_layout(G)

# 转换时间序列数据为帧
frame_interval = 5  # ms
num_frames = int(duration/ms/frame_interval)
time_points = np.linspace(0, duration/ms, num_frames)

def update(frame):
    plt.clf()
    current_time = time_points[frame]
    
    # 1. 膜电位随时间变化（左上）
    plt.subplot(221)
    time_mask = t_points <= current_time
    for i in range(0, N, 10):
        plt.plot(t_points[time_mask], v_points[i][time_mask])
    plt.xlabel('Time (ms)')
    plt.ylabel('v (mV)')
    plt.title('Membrane Potential over Time')
    plt.ylim(-1, 16)
    
    # 2. 脉冲发放图（右上）
    plt.subplot(222)
    spike_mask = spike_times <= current_time
    plt.plot(spike_times[spike_mask], spike_indices[spike_mask], '.k')
    plt.xlabel('Time (ms)')
    plt.ylabel('Neuron index')
    plt.title('Spike Raster Plot')
    plt.xlim(0, duration/ms)
    plt.ylim(-1, N)
    
    # 3. 网络状态（左下）
    plt.subplot(223)
    time_idx = np.searchsorted(t_points, current_time)
    if time_idx >= len(t_points):
        time_idx = len(t_points) - 1
    current_v = v_points.T[time_idx]
    node_colors = plt.cm.viridis(current_v/15)  # 根据膜电位着色
    nx.draw(G, pos, node_size=20, node_color=node_colors,
            with_labels=False, arrows=False)
    plt.title('Network State')
    
    # 4. 累积发放率（右下）
    plt.subplot(224)
    current_spikes = spike_indices[spike_times <= current_time]
    spike_counts = np.histogram(current_spikes, bins=N, range=(0, N))[0]
    plt.bar(range(N), spike_counts)
    plt.xlabel('Neuron index')
    plt.ylabel('Spike count')
    plt.title('Cumulative Firing Rate')
    
    plt.tight_layout()

# 创建动画
anim = FuncAnimation(fig, update, frames=num_frames, 
                    interval=50, blit=False)

# 保存动画
writer = animation.PillowWriter(fps=20)
anim.save('snn_animation.gif', writer=writer)
plt.close()

print("动画已保存为 'snn_animation.gif'") 